"""
合同要素抽取
"""
import json
from typing import Any, Dict
from core.milvus_vector import MilvusVector
from core.config import settings
from core.reranker import ReRanker
from loguru import logger


async def recall(fileId: str, rule: Dict[str, Any]) -> Dict:
    """
    合同要素的内容抽取
    """
    info = {'fileId': fileId, 'rule': rule}
    logger.info(f'文本召回: {json.dumps(info, ensure_ascii=False)}')

    # 召回要素相关的片段
    milvus = MilvusVector()
    rule_name = rule['name']
    extract_hint = rule['extractHint'] or ''

    query_parts = []
    if extract_hint != '':
        query_parts.append(f'按照如下指示：`{extract_hint}`')
    query_parts.append(f'找出`{rule_name}`所指向的文本内容')
    query = '，'.join(query_parts) + "。"

    documents = milvus.hybrid_search(fileId, query, top_k=30)

    # 文档重排
    reranker = ReRanker()
    documents = reranker.rerank(query, documents, settings.TOP_K)

    if settings.DEBUG:
        logger.debug('召回的文本片段：')
        for doc in documents:
            logger.debug(doc)

    # 按合同原始顺序拼接成markdown形式的文本, 标题也需整合进去
    retrieved = dict()
    paragraph_dict = dict()
    for doc in documents:
        metadata = doc.metadata

        # 加入两层title
        top_toc_index = metadata.get('top_toc_index', None)
        if top_toc_index is not None:
            retrieved[top_toc_index] = metadata.get('top_toc_title', '')
        top_toc_id = metadata.get('top_toc_id', None)
        if top_toc_id is not None:
            paragraph_dict[top_toc_id] = metadata.get('top_toc_title', '')

        sub_toc_index = metadata.get('sub_toc_index', None)
        if sub_toc_index is not None:
            retrieved[sub_toc_index] = metadata.get('sub_toc_title', '')
        sub_toc_id = metadata.get('sub_toc_id', None)
        if sub_toc_id is not None:
            paragraph_dict[sub_toc_id] = metadata.get('sub_toc_id', '')

        # 适当扩展召回的上下文
        previous_index = metadata.get('previous_index', None)
        if previous_index is not None:
            retrieved[previous_index] = metadata.get('previous_content', '')
        previous_id = metadata.get('previous_id', None)
        if previous_id is not None:
            paragraph_dict[previous_id] = metadata.get('previous_content', '')

        next_index = metadata.get('next_index', None)
        if next_index is not None:
            retrieved[next_index] = metadata.get('next_content', '')
        next_id = metadata.get('next_id', None)
        if next_id is not None:
            paragraph_dict[next_id] = metadata.get('next_content', '')

        next2_index = metadata.get('next2_index', None)
        if next2_index is not None:
            retrieved[next2_index] = metadata.get('next2_content', '')
        next2_id = metadata.get('next2_id', None)
        if next2_id is not None:
            paragraph_dict[next2_id] = metadata.get('next2_content', '')

        retrieved[metadata['block_index']] = metadata['block_content']
        if metadata['id'] is not None:
            paragraph_dict[metadata['id']] = doc.content

    retrieved = sorted(retrieved.items(), key=lambda x: x[0])
    related_texts = [content for _, content in retrieved]

    # 返回结果
    extract_result = {
        'fileId': fileId,             # 合同文件id
        'id': rule.get('id', None),   # 要素id
        'name': rule_name,            # 要素名称
        'extractHint': extract_hint,  # 要素抽取的hint
        'related_ids': list(paragraph_dict.keys()), # 召回的相关段落id
        'related_texts': '\n'.join(related_texts)   # 召回的相关文本
    }
    logger.info(f'召回结果: {json.dumps(extract_result, ensure_ascii=False)}')
    return extract_result
